import math, sys, os, io, subprocess
from PIL import Image, ImageChops
from functools import partial
from PySide6.QtWidgets import QApplication, QMainWindow, QFileDialog
from PySide6.QtCore import QFile, Signal, QObject, QThreadPool, QRunnable
from PySide6.QtUiTools import QUiLoader

class PNGquantSignal(QObject):
    print_text = Signal(str)  # 信息显示信号

class PNGquant(QRunnable):
    def __init__(self, params, max_width, max_height, inself):
        super().__init__()
        self.png_file = params
        self.max_width = max_width
        self.max_height = max_height
        self.inself = inself
        self.signals_info = PNGquantSignal()

        # 连接信号到方法
        self.signals_info.print_text.connect(self.inself.log_info)

    def run(self):
        """调整图片大小并压缩。"""
        for png_file in self.png_file:
            inpu_file, out_file= png_file
            with Image.open(inpu_file) as img:
                checked_data = self.inself.window.crop_whitespace.isChecked()
                if checked_data:
                    # 从图片中心开始裁剪
                    left = (img.width - self.max_width) / 2
                    top = (img.height - self.max_height) / 2
                    right = (img.width + self.max_width) / 2
                    bottom = (img.height + self.max_height) / 2

                    # 确保裁剪坐标不越界
                    left = max(0, left)
                    top = max(0, top)
                    right = min(img.width, right)
                    bottom = min(img.height, bottom)

                    # 裁剪图片
                    new_img = img.crop((left, top, right, bottom))
                else:
                    new_img = img

                if self.inself.window.img_resized.text() != 100:
                    width, height = new_img.size
                    width = round(int(self.inself.window.img_resized.text()) / 100 * width)
                    height = round(int(self.inself.window.img_resized.text()) / 100 * height)
                    new_img = new_img.resize((width, height), resample=Image.LANCZOS)
                else:
                    new_img = img

                img_byte_array = io.BytesIO()
                new_img.save(img_byte_array, format='PNG')
                img_byte_array.seek(0)

                cmd = f"{os.path.join(os.path.dirname(os.path.abspath(__file__)), pngquant_exe)} --force --verbose --speed=1 --ordered 256 - --output {out_file}"
                proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
                stdout, stderr = proc.communicate(input=img_byte_array.read())
                if proc.returncode != 0:
                    try:
                        # 假设外部程序输出使用GBK编码
                        error_message = stderr.decode('GBK')
                    except UnicodeDecodeError as e:
                        # 如果解码失败，尝试忽略解码错误
                        error_message = stderr.decode('GBK', errors='ignore')

                    self.signals_info.print_text.emit(
                        f"压缩 {inpu_file} 出错：返回码 {proc.returncode}, 错误信息：{error_message}")

                elif proc.returncode == 99:
                    self.signals_info.print_text.emit(f"{inpu_file} 未处理。")

                else:
                    self.signals_info.print_text.emit(f"{inpu_file} 压缩成功。")

    def trim(self, im):
        """从图像中去除透明边缘，并返回裁剪后的图像。"""
        bg = Image.new(im.mode, im.size, im.getpixel((0, 0)))
        diff = ImageChops.difference(im, bg)
        bbox = diff.getbbox()
        if bbox:
            return im.crop(bbox)
        return im

class Worker(QObject):
    finished_signal = Signal()  # 任务完成信号

    def __init__(self, file_pairs, file_list, in_self):
        super().__init__()
        self.params = file_pairs
        self.inself = in_self
        self.max_width, self.max_height = self.find_max_dimensions(file_list)

    def start_processing(self):
        # 分割文件列表到多个子列表，每个子列表一个线程处理
        num_threads = min(len(self.params), os.cpu_count())
        chunk_size = math.ceil(len(self.params) / num_threads)
        file_chunks = [self.params[i * chunk_size:(i + 1) * chunk_size] for i in range(num_threads)]
        for chunk in file_chunks:
            if chunk:
                PNG_worker = PNGquant(chunk, self.max_width, self.max_height, self.inself)
                QThreadPool.globalInstance().start(PNG_worker)
        # 等待所有线程完成
        QThreadPool.globalInstance().waitForDone()
        self.finished_signal.emit()

    def find_max_dimensions(self, file_list):
        """计算提供的文件列表中所有图像的最大宽度和高度。"""
        max_width, max_height = 0, 0
        for file_path in file_list:
            with Image.open(file_path) as img:
                trimmed_img = self.trim(img)
                width, height = trimmed_img.size
                max_width = max(max_width, width)
                max_height = max(max_height, height)
        if max_width % 2 != 0:
            max_width += 1

        if max_height % 2 !=0:
            max_height += 1
        return max_width, max_height

    def trim(self, im):
        """裁剪图像，移除所有透明边缘。"""
        bg = Image.new(im.mode, im.size, im.getpixel((0, 0)))
        diff = ImageChops.difference(im, bg)
        bbox = diff.getbbox()
        if bbox:
            return im.crop(bbox)
        return im


class MyMainWindow(QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        ui_file = QFile("ui/pngquant.ui")
        ui_file.open(QFile.ReadOnly)
        self.window = loader.load(ui_file)
        ui_file.close()

        self.setCentralWidget(self.window)
        self.setWindowTitle("PNGquant v1.0.0")
        self.window.ExporButton.clicked.connect(self.ExporButton_on_off)  # 压缩
        self.window.Button_projects.clicked.connect(partial(self.select_directory, "项目"))  # 选择项目目录按钮
        self.window.Button_expor.clicked.connect(partial(self.select_directory, "导出"))  # 导出目录按钮
        self.window.save_project_dir.stateChanged.connect(self.save_project_directory)  # 导出到项目目录
        self.window.crop_whitespace.setChecked(True)  # 默认裁剪
        self.window.one_crop.setChecked(False)  # 独裁剪（非序列）
        self.window.one_crop.setEnabled(False)
        self.task_finished = 0

    # 设置导出路径
    def save_project_directory(self):
        global output
        global out_path
        set_save_path = self.window.save_project_dir.isChecked()
        if set_save_path == True:
            out_path = self.window.output_folder.text()  # 导出目录
            self.window.output_folder.setEnabled(False)
            self.window.Button_expor.setEnabled(False)
            self.window.output_folder.setText(".\\压缩")

        else:
            self.window.output_folder.setEnabled(True)
            self.window.Button_expor.setEnabled(True)
            self.window.output_folder.setText(out_path)
            output = out_path

    # 选择目录
    def select_directory(self, Text):
        if str(Text) in "项目":
            str_text = "选择项目目录"
        elif str(Text) in "导出":
            str_text = "选择导出目录"
        selected_dir = QFileDialog.getExistingDirectory(None, str(str_text), "C:/")
        selected_dir = selected_dir.replace("/", "\\")  # 路径\改变
        if selected_dir:
            if str(Text) in "项目":
                self.window.project.setText(str(selected_dir))  # 项目目录
            elif str(Text) in "导出":
                self.window.output_folder.setText(str(selected_dir))  # 导出目录

    def ExporButton_on_off(self):
        """处理导出按钮的点击事件，控制压缩开始或停止。"""
        if self.window.ExporButton.text() == "停止":
            self.window.ExporButton.setText('压缩')
        else:
            self.window.progressBar.setValue(0)
            self.window.ExporButton.setText('停止')
            self.window.plainTextEdit.clear()
            self.start_export()

    def start_export(self):
        """开始导出处理，读取UI中的路径并启动文件处理。"""
        project_path = self.window.project.text()  # 项目目录
        set_save_path = self.window.output_folder.text()  # 导出目录
        set_dir = self.window.save_project_dir.isChecked()
        if set_dir:
            set_save_path = project_path + "\\压缩"

        # 如果导出目录不存在就创建目录
        if not os.path.exists(set_save_path):
            os.makedirs(set_save_path)

        self.find_file(project_path, set_save_path)

    def find_file(self, src_folder, dst_folder):
        """查找目标文件夹中的所有PNG文件并启动压缩过程。"""
        file_list = [os.path.join(root, file) for root, dirs, files in os.walk(src_folder) for file in files if file.lower().endswith('.png')]
        self.window.progressBar.setMaximum(len(file_list))
        file_pairs = [(file, file.replace(src_folder, dst_folder)) for file in file_list]
        worker = Worker(file_pairs, file_list, self)
        worker.finished_signal.connect(self.worker_finished)
        worker.start_processing()

    def worker_finished(self):
        """处理Worker完成信号，更新UI状态。"""
        self.window.ExporButton.setText('压缩')

    # 信息窗口
    def log_info(self, message):
        self.task_finished += 1
        self.window.progressBar.setValue(self.task_finished)
        self.window.plainTextEdit.append(message)



if __name__ == '__main__':
    loader = QUiLoader()  # 先初始化，可以是pyside6.6.1的版本问题
    app = QApplication(sys.argv)
    exe_path = sys.executable
    exe_dir = os.path.dirname(exe_path)
    pngquant_exe = str(exe_dir) + "\\pngquant\\pngquant.exe"
    window = MyMainWindow()
    window.show()
    sys.exit(app.exec())